{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'PillowWriter'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-4-51925c8ee522>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     18\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mmatplotlib\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mpyplot\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mplt\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     19\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mscipy\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mintegrate\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mintegrate\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 20\u001b[1;33m \u001b[1;32mimport\u001b[0m \u001b[0mmatplotlib\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0manimation\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0manimation\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mPillowWriter\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     21\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     22\u001b[0m \u001b[1;32mclass\u001b[0m \u001b[0mDoublePendulum\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'PillowWriter'"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "General Numerical Solver for the 1D Time-Dependent Schrodinger's equation.\n",
    "\n",
    "adapted from code at http://matplotlib.sourceforge.net/examples/animation/double_pendulum_animated.py\n",
    "\n",
    "Double pendulum formula translated from the C code at\n",
    "http://www.physics.usyd.edu.au/~wheat/dpend_html/solve_dpend.c\n",
    "\n",
    "author: Jake Vanderplas\n",
    "email: vanderplas@astro.washington.edu\n",
    "website: http://jakevdp.github.com\n",
    "license: BSD\n",
    "Please feel free to use and modify this, but keep the above information. Thanks!\n",
    "\"\"\"\n",
    "\n",
    "from numpy import sin, cos\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scipy.integrate as integrate\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "class DoublePendulum:\n",
    "    \"\"\"Double Pendulum Class\n",
    "\n",
    "    init_state is [theta1, omega1, theta2, omega2] in degrees,\n",
    "    where theta1, omega1 is the angular position and velocity of the first\n",
    "    pendulum arm, and theta2, omega2 is that of the second pendulum arm\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 init_state = [120, 0, -20, 0],\n",
    "                 L1=1.0,  # length of pendulum 1 in m\n",
    "                 L2=1.0,  # length of pendulum 2 in m\n",
    "                 M1=1.0,  # mass of pendulum 1 in kg\n",
    "                 M2=1.0,  # mass of pendulum 2 in kg\n",
    "                 G=9.8,  # acceleration due to gravity, in m/s^2\n",
    "                 origin=(0, 0)): \n",
    "        self.init_state = np.asarray(init_state, dtype='float')\n",
    "        self.params = (L1, L2, M1, M2, G)\n",
    "        self.origin = origin\n",
    "        self.time_elapsed = 0\n",
    "\n",
    "        self.state = self.init_state * np.pi / 180.\n",
    "    \n",
    "    def position(self):\n",
    "        \"\"\"compute the current x,y positions of the pendulum arms\"\"\"\n",
    "        (L1, L2, M1, M2, G) = self.params\n",
    "\n",
    "        x = np.cumsum([self.origin[0],\n",
    "                       L1 * sin(self.state[0]),\n",
    "                       L2 * sin(self.state[2])])\n",
    "        y = np.cumsum([self.origin[1],\n",
    "                       -L1 * cos(self.state[0]),\n",
    "                       -L2 * cos(self.state[2])])\n",
    "        return (x, y)\n",
    "\n",
    "    def energy(self):\n",
    "        \"\"\"compute the energy of the current state\"\"\"\n",
    "        (L1, L2, M1, M2, G) = self.params\n",
    "\n",
    "        x = np.cumsum([L1 * sin(self.state[0]),\n",
    "                       L2 * sin(self.state[2])])\n",
    "        y = np.cumsum([-L1 * cos(self.state[0]),\n",
    "                       -L2 * cos(self.state[2])])\n",
    "        vx = np.cumsum([L1 * self.state[1] * cos(self.state[0]),\n",
    "                        L2 * self.state[3] * cos(self.state[2])])\n",
    "        vy = np.cumsum([L1 * self.state[1] * sin(self.state[0]),\n",
    "                        L2 * self.state[3] * sin(self.state[2])])\n",
    "\n",
    "        U = G * (M1 * y[0] + M2 * y[1])\n",
    "        K = 0.5 * (M1 * np.dot(vx, vx) + M2 * np.dot(vy, vy))\n",
    "\n",
    "        return U + K\n",
    "\n",
    "    def dstate_dt(self, state, t):\n",
    "        \"\"\"compute the derivative of the given state\"\"\"\n",
    "        (M1, M2, L1, L2, G) = self.params\n",
    "\n",
    "        dydx = np.zeros_like(state)\n",
    "        dydx[0] = state[1]\n",
    "        dydx[2] = state[3]\n",
    "\n",
    "        cos_delta = cos(state[2] - state[0])\n",
    "        sin_delta = sin(state[2] - state[0])\n",
    "\n",
    "        den1 = (M1 + M2) * L1 - M2 * L1 * cos_delta * cos_delta\n",
    "        dydx[1] = (M2 * L1 * state[1] * state[1] * sin_delta * cos_delta\n",
    "                   + M2 * G * sin(state[2]) * cos_delta\n",
    "                   + M2 * L2 * state[3] * state[3] * sin_delta\n",
    "                   - (M1 + M2) * G * sin(state[0])) / den1\n",
    "\n",
    "        den2 = (L2 / L1) * den1\n",
    "        dydx[3] = (-M2 * L2 * state[3] * state[3] * sin_delta * cos_delta\n",
    "                   + (M1 + M2) * G * sin(state[0]) * cos_delta\n",
    "                   - (M1 + M2) * L1 * state[1] * state[1] * sin_delta\n",
    "                   - (M1 + M2) * G * sin(state[2])) / den2\n",
    "        \n",
    "        return dydx\n",
    "\n",
    "    def step(self, dt):\n",
    "        \"\"\"execute one time step of length dt and update state\"\"\"\n",
    "        self.state = integrate.odeint(self.dstate_dt, self.state, [0, dt])[1]\n",
    "        self.time_elapsed += dt\n",
    "\n",
    "#------------------------------------------------------------\n",
    "# set up initial state and global variables\n",
    "pendulum = DoublePendulum([180., 0.0, -20., 0.0])\n",
    "dt = 1./30 # 30 fps\n",
    "\n",
    "#------------------------------------------------------------\n",
    "# set up figure and animation\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, aspect='equal', autoscale_on=False,\n",
    "                     xlim=(-2, 2), ylim=(-2, 2))\n",
    "ax.grid()\n",
    "\n",
    "line, = ax.plot([], [], 'o-', lw=2)\n",
    "time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)\n",
    "energy_text = ax.text(0.02, 0.90, '', transform=ax.transAxes)\n",
    "\n",
    "def init():\n",
    "    \"\"\"initialize animation\"\"\"\n",
    "    line.set_data([], [])\n",
    "    time_text.set_text('')\n",
    "    energy_text.set_text('')\n",
    "    return line, time_text, energy_text\n",
    "\n",
    "def animate(i):\n",
    "    \"\"\"perform animation step\"\"\"\n",
    "    global pendulum, dt\n",
    "    pendulum.step(dt)\n",
    "    \n",
    "    line.set_data(*pendulum.position())\n",
    "    time_text.set_text('time = %.1f' % pendulum.time_elapsed)\n",
    "    energy_text.set_text('energy = %.3f J' % pendulum.energy())\n",
    "    return line, time_text, energy_text\n",
    "\n",
    "# choose the interval based on dt and the time to animate one step\n",
    "from time import time\n",
    "t0 = time()\n",
    "animate(0)\n",
    "t1 = time()\n",
    "interval = 1000 * dt - (t1 - t0)\n",
    "\n",
    "ani = animation.FuncAnimation(fig, animate, frames=300,\n",
    "                              interval=interval, blit=True, init_func=init)\n",
    "\n",
    "# save the animation as an mp4.  This requires ffmpeg or mencoder to be\n",
    "# installed.  The extra_args ensure that the x264 codec is used, so that\n",
    "# the video can be embedded in html5.  You may need to adjust this for\n",
    "# your system: for more information, see\n",
    "# http://matplotlib.sourceforge.net/api/animation_api.html\n",
    "ani.save('double_pendulum.mp4', fps=30, writer='pillow')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "b3ba2566441a7c06988d0923437866b63cedc61552a5af99d1f4fb67d367b25f"
  },
  "kernelspec": {
   "display_name": "Python 3.8.3 64-bit ('base': conda)",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
